import math
import torch
import numpy as np

def kl_gaussian(mu, log_var):
    # mu_size:[BatchSize,LatentDim], log_var_size:[BatchSize,LatentDim]
    # 比较当前潜在编码对应分布与标准正态分布之间的KL散度
    return -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp(), dim=1)

def gaussian_sample(mu, log_var):
    # mu_size:[BatchSize,LatentDim], log_var_size:[BatchSize,LatentDim]
    # log_var:方差的对数
    # 从潜在编码对应的连续分布进行随机采样
    gaussian_noise = torch.randn(mu.size())
    if mu.is_cuda:
        gaussian_noise = gaussian_noise.cuda()
    return mu + torch.exp(log_var * 0.5) * gaussian_noise

def normal_distribution(x, mean, sigma):
    # x_size:[BatchSize,1],mean:float,sigma:float
    # 计算张量x的各个元素的概率密度函数值
    return torch.exp(-1*((x-mean)**2)/(2*sigma**2))/(np.sqrt(2*np.pi) * sigma)

def mixed_2_gaussian_distribution(x, mean1, sigma1, mean2, sigma2):
    # x_size:[BatchSize,1],mean1/2:float,sigma1/2:float
    # 计算两个高斯密度函数混合而成的高斯混合密度
    return 0.5*normal_distribution(x,mean1,sigma1)+0.5*normal_distribution(x,mean2,sigma2)

def get_poisson_distribution(lam,k):
    # lam:float,k:integer
    # 计算泊松分布期望为lam时，变量值取k的概率
    return math.pow(lam,k)/math.factorial(k)*math.exp(-1*lam)

def get_poisson_normalizer(lam,lower=1,upper=100):
    # lam:float,lower:integer,upper:integer
    # 计算泊松分布截断在区间[lower,upper]上时的normalization因子，即该区间上概率之和
    poisson_probs=[get_poisson_distribution(lam,x) for x in range(lower,upper+1)]
    return np.sum(poisson_probs)

def get_normalized_poisson(lam,lower,upper):
    # lam:float,lower:integer,upper:integer
    # 计算经过normalization的泊松分布
    probs=[get_poisson_distribution(lam,k) for k in range(lower,upper+1)]
    normalizer=sum(probs)
    probs=list(np.array(probs)/normalizer)
    return probs

def stretched_sigmoid(x,mean=0,k=1):
    # x:逻辑上应该是integer, 当前step数
    # 通过拉伸sigmoid获取一个近似的阶跃函数
    return 1/(1+np.exp(-1*k*(x-mean)))

def stretched_sigmoid_torch(x,mean=0,k=1):
    # x:应当是tensor
    # 通过拉伸Sigmoid获取一个近似的阶跃函数
    return 1/(1+torch.exp(-1*k*(x-mean)))